#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
点云图像融合Web API后端
基于Flask框架，提供RESTful API接口
"""

import os
import io
import json
import numpy as np
import cv2
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
from werkzeug.utils import secure_filename
import tempfile
import base64

# 导入Web适配的融合处理函数
import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from fusion_web import (
    read_bin, read_calib, get_fov_mask,
    lidar2image, perform_fusion_web
)

app = Flask(__name__)
CORS(app)  # 允许跨域请求

# 配置
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024  # 100MB最大文件大小
UPLOAD_FOLDER = 'temp_uploads'
ALLOWED_EXTENSIONS = {
    'pointcloud': {'bin'},
    'image': {'png', 'jpg', 'jpeg'},
    'calib': {'txt'}
}

# 确保上传文件夹存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

def allowed_file(filename, file_type):
    """检查文件扩展名是否允许"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS[file_type]

def save_uploaded_file(file, file_type):
    """保存上传的文件并返回文件路径"""
    if file and allowed_file(file.filename, file_type):
        filename = secure_filename(file.filename)
        # 添加时间戳避免文件名冲突
        import time
        timestamp = str(int(time.time()))
        filename = f"{timestamp}_{filename}"
        filepath = os.path.join(UPLOAD_FOLDER, filename)
        file.save(filepath)
        return filepath
    return None

@app.route('/api/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return jsonify({
        'status': 'ok',
        'message': '点云图像融合API服务正常运行'
    })

@app.route('/api/upload', methods=['POST'])
def upload_files():
    """文件上传接口"""
    try:
        # 检查是否有文件上传
        if 'pointcloud' not in request.files or \
           'image' not in request.files or \
           'calib' not in request.files:
            return jsonify({
                'success': False,
                'error': '请上传所有必需的文件（点云、图像、标定文件）'
            }), 400

        files = {
            'pointcloud': request.files['pointcloud'],
            'image': request.files['image'],
            'calib': request.files['calib']
        }

        # 保存文件
        file_paths = {}
        for file_type, file in files.items():
            if file.filename == '':
                return jsonify({
                    'success': False,
                    'error': f'{file_type}文件不能为空'
                }), 400
            
            filepath = save_uploaded_file(file, file_type)
            if filepath is None:
                return jsonify({
                    'success': False,
                    'error': f'{file_type}文件格式不支持'
                }), 400
            
            file_paths[file_type] = filepath

        return jsonify({
            'success': True,
            'message': '文件上传成功',
            'file_paths': file_paths
        })

    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'文件上传失败: {str(e)}'
        }), 500

@app.route('/api/process', methods=['POST'])
def process_fusion():
    """点云图像融合处理接口"""
    try:
        data = request.get_json()
        
        if not data or 'file_paths' not in data:
            return jsonify({
                'success': False,
                'error': '缺少文件路径信息'
            }), 400

        file_paths = data['file_paths']
        
        # 验证文件路径
        required_files = ['pointcloud', 'image', 'calib']
        for file_type in required_files:
            if file_type not in file_paths:
                return jsonify({
                    'success': False,
                    'error': f'缺少{file_type}文件路径'
                }), 400
            
            if not os.path.exists(file_paths[file_type]):
                return jsonify({
                    'success': False,
                    'error': f'{file_type}文件不存在'
                }), 400

        # 执行融合处理
        result = perform_fusion_web(
            file_paths['pointcloud'],
            file_paths['image'],
            file_paths['calib']
        )

        # 清理临时文件
        for filepath in file_paths.values():
            try:
                if os.path.exists(filepath):
                    os.remove(filepath)
                    print(f'✅ 已清理临时文件: {filepath}')
            except Exception as e:
                print(f'⚠️ 无法清理文件 {filepath}: {str(e)}')

        return jsonify({
            'success': True,
            'message': '融合处理完成',
            'result': result
        })

    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'融合处理失败: {str(e)}'
        }), 500

# perform_fusion函数已移至fusion_web.py中的perform_fusion_web函数

@app.route('/api/visualize', methods=['POST'])
def generate_visualization():
    """生成可视化数据接口"""
    try:
        data = request.get_json()
        
        if not data or 'points' not in data or 'colors' not in data:
            return jsonify({
                'success': False,
                'error': '缺少点云或颜色数据'
            }), 400

        points = np.array(data['points'])
        colors = np.array(data['colors'])
        
        # 生成Three.js格式的数据
        visualization_data = {
            'positions': points.flatten().tolist(),
            'colors': (colors / 255.0).flatten().tolist(),  # 归一化到0-1
            'point_count': len(points),
            'bounds': {
                'x': [float(points[:, 0].min()), float(points[:, 0].max())],
                'y': [float(points[:, 1].min()), float(points[:, 1].max())],
                'z': [float(points[:, 2].min()), float(points[:, 2].max())]
            }
        }
        
        return jsonify({
            'success': True,
            'visualization_data': visualization_data
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'生成可视化数据失败: {str(e)}'
        }), 500

@app.route('/api/export', methods=['POST'])
def export_data():
    """导出处理结果接口 - 增强版"""
    try:
        data = request.get_json()

        if not data or 'points' not in data or 'colors' not in data:
            return jsonify({
                'success': False,
                'error': '缺少导出数据'
            }), 400

        points = np.array(data['points'])
        colors = np.array(data['colors'])
        include_intensity = data.get('include_intensity', True)
        export_format = data.get('format', 'ply').lower()
        filename = data.get('filename', 'colored_pointcloud')

        # 验证文件名
        import re
        if not re.match(r'^[a-zA-Z0-9_\-\u4e00-\u9fa5]+$', filename):
            return jsonify({
                'success': False,
                'error': '文件名只能包含字母、数字、下划线、连字符和中文字符'
            }), 400

        # 生成文件内容
        if export_format == 'ply':
            content, mimetype = generate_ply_content(points, colors, include_intensity)
        elif export_format == 'pcd':
            content, mimetype = generate_pcd_content(points, colors, include_intensity)
        elif export_format == 'xyz':
            content, mimetype = generate_xyz_content(points, colors, include_intensity)
        elif export_format == 'txt':
            content, mimetype = generate_txt_content(points, colors, include_intensity)
        else:
            return jsonify({
                'success': False,
                'error': f'不支持的导出格式: {export_format}'
            }), 400

        # 创建临时文件
        with tempfile.NamedTemporaryFile(mode='w', suffix=f'.{export_format}', delete=False, encoding='utf-8') as f:
            f.write(content)
            temp_path = f.name

        return send_file(
            temp_path,
            as_attachment=True,
            download_name=f'{filename}.{export_format}',
            mimetype=mimetype
        )

    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'导出数据失败: {str(e)}'
        }), 500

def generate_ply_content(points, colors, include_intensity):
    """生成PLY格式内容"""
    content = "ply\n"
    content += "format ascii 1.0\n"
    content += f"element vertex {len(points)}\n"
    content += "property float x\n"
    content += "property float y\n"
    content += "property float z\n"
    content += "property uchar red\n"
    content += "property uchar green\n"
    content += "property uchar blue\n"

    if include_intensity:
        content += "property float intensity\n"

    content += "end_header\n"

    for i in range(len(points)):
        point = points[i]
        color = colors[i]

        if include_intensity and len(color) >= 4:
            # 格式: x y z r g b intensity (intensity在color数组的第4位)
            content += f"{point[0]} {point[1]} {point[2]} {int(color[0])} {int(color[1])} {int(color[2])} {color[3]}\n"
        else:
            content += f"{point[0]} {point[1]} {point[2]} {int(color[0])} {int(color[1])} {int(color[2])}\n"

    return content, 'application/octet-stream'

def generate_pcd_content(points, colors, include_intensity):
    """生成PCD格式内容"""
    fields = "x y z rgb intensity" if include_intensity else "x y z rgb"
    size = "4 4 4 4 4" if include_intensity else "4 4 4 4"
    type_str = "F F F U F" if include_intensity else "F F F U"
    count = "1 1 1 1 1" if include_intensity else "1 1 1 1"

    content = "# .PCD v0.7 - Point Cloud Data file format\n"
    content += "VERSION 0.7\n"
    content += f"FIELDS {fields}\n"
    content += f"SIZE {size}\n"
    content += f"TYPE {type_str}\n"
    content += f"COUNT {count}\n"
    content += f"WIDTH {len(points)}\n"
    content += "HEIGHT 1\n"
    content += "VIEWPOINT 0 0 0 1 0 0 0\n"
    content += f"POINTS {len(points)}\n"
    content += "DATA ascii\n"

    for i in range(len(points)):
        point = points[i]
        color = colors[i]
        rgb = (int(color[0]) << 16) | (int(color[1]) << 8) | int(color[2])

        if include_intensity and len(color) >= 4:
            # 格式: x y z rgb intensity (intensity在color数组的第4位)
            content += f"{point[0]} {point[1]} {point[2]} {rgb} {color[3]}\n"
        else:
            content += f"{point[0]} {point[1]} {point[2]} {rgb}\n"

    return content, 'application/octet-stream'

def generate_xyz_content(points, colors, include_intensity):
    """生成XYZ格式内容"""
    content = ""

    for i in range(len(points)):
        point = points[i]
        color = colors[i]

        if include_intensity and len(color) >= 4:
            # 格式: x y z r g b intensity (intensity在color数组的第4位)
            content += f"{point[0]} {point[1]} {point[2]} {int(color[0])} {int(color[1])} {int(color[2])} {color[3]}\n"
        else:
            # 格式: x y z r g b
            content += f"{point[0]} {point[1]} {point[2]} {int(color[0])} {int(color[1])} {int(color[2])}\n"

    return content, 'text/plain'

def generate_txt_content(points, colors, include_intensity):
    """生成TXT格式内容"""
    header = "# Point Cloud Data\n"
    header += f"# Format: {'x y z r g b intensity' if include_intensity else 'x y z r g b'}\n"
    header += f"# Points: {len(points)}\n"
    header += "#\n"

    content, _ = generate_xyz_content(points, colors, include_intensity)
    return header + content, 'text/plain'

@app.errorhandler(413)
def too_large(e):
    """文件过大错误处理"""
    return jsonify({
        'success': False,
        'error': '上传文件过大，请确保文件小于100MB'
    }), 413

@app.errorhandler(404)
def not_found(e):
    """404错误处理"""
    return jsonify({
        'success': False,
        'error': 'API接口不存在'
    }), 404

@app.errorhandler(500)
def internal_error(e):
    """500错误处理"""
    return jsonify({
        'success': False,
        'error': '服务器内部错误'
    }), 500

if __name__ == '__main__':
    print("启动点云图像融合Web API服务...")
    print("API文档:")
    print("  GET  /api/health     - 健康检查")
    print("  POST /api/upload     - 文件上传")
    print("  POST /api/process    - 融合处理")
    print("  POST /api/visualize  - 生成可视化")
    print("  POST /api/export     - 导出结果")
    print("\n服务地址: http://localhost:5000")
    
    app.run(
        host='0.0.0.0',
        port=5000,
        debug=True,
        threaded=True
    )
